{
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "import jax\n",
        "jax.__version__"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "vcRv42-IQfZc",
        "outputId": "c3acc10d-022e-49ed-b257-da39503bee03"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'0.4.14'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 1
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZE6H-14Adw5d"
      },
      "source": [
        "## DeviceArray properties"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "Xw26fWWnf8eB"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import jax.numpy as jnp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tyk_72Oxq4WG",
        "outputId": "dbd7570a-f26a-4846-84da-546455e33792"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([    1,    42, 31337])"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ],
      "source": [
        "np.array([1, 42, 31337])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pZvJaXWTBm04",
        "outputId": "9815ad1e-2231-4358-b9c6-cdc31f7101f2"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([    1,    42, 31337], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ],
      "source": [
        "jnp.array([1, 42, 31337])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JF5SSt3DBskh",
        "outputId": "5725bcf7-3b5f-4cad-aefe-ea5d1991011c"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "31380"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ],
      "source": [
        "np.sum([1, 42, 31337])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5O9F4CNYBymQ",
        "outputId": "ad00b7ee-b32e-41b7-eae6-b4bdc92b574a"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "sum requires ndarray or scalar arguments, got <class 'list'> at position 0.\n"
          ]
        }
      ],
      "source": [
        "try:\n",
        "  jnp.sum([1, 42, 31337])\n",
        "except TypeError as e:\n",
        "  print(e)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "x76jmBRaHvQp",
        "outputId": "1e2e8231-6e1c-4814-95a9-cd4e78343eae"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(31380, dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ],
      "source": [
        "jnp.sum(jnp.array([1, 42, 31337]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "sX1UevU0K5AW"
      },
      "outputs": [],
      "source": [
        "arr = jnp.array([1, 42, 31337])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pDSN47PTRZHG",
        "outputId": "68594806-9c95-4870-ad93-817cf1302330"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "1"
            ]
          },
          "metadata": {},
          "execution_count": 9
        }
      ],
      "source": [
        "arr.ndim"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6pghF6HXLx5O",
        "outputId": "351b355a-4346-4649-fbee-a69d6fc3e5e1"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(3,)"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ],
      "source": [
        "arr.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AXaNW2zaSzuR",
        "outputId": "dda018c2-f981-482f-bb12-0c3119eb17f1"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "dtype('int32')"
            ]
          },
          "metadata": {},
          "execution_count": 11
        }
      ],
      "source": [
        "arr.dtype"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "G2GJkI3uRYYC",
        "outputId": "55a064a5-9d87-4283-dc08-889404579800"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "3"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ],
      "source": [
        "arr.size"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rDB_80JCSv_b",
        "outputId": "8055e6d3-e3f5-43f1-aa9f-eef3abe2b5e9"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "12"
            ]
          },
          "metadata": {},
          "execution_count": 13
        }
      ],
      "source": [
        "arr.nbytes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BXbIZE5Vd39Z"
      },
      "source": [
        "## Devices"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "1GAG9wYcaoSL"
      },
      "outputs": [],
      "source": [
        "import jax"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gInETIcSakx6",
        "outputId": "253f0974-1d28-43c0-b7cb-0929d09f9089"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[gpu(id=0)]"
            ]
          },
          "metadata": {},
          "execution_count": 15
        }
      ],
      "source": [
        "jax.devices()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AUOges5na8X9",
        "outputId": "3a2208c7-8125-4940-a5a1-41e6a9a101c1"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[CpuDevice(id=0)]"
            ]
          },
          "metadata": {},
          "execution_count": 16
        }
      ],
      "source": [
        "jax.devices('cpu')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cXKsGEZObOkH",
        "outputId": "6eb84e74-300e-4425-d67f-669f75230974"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "1"
            ]
          },
          "metadata": {},
          "execution_count": 17
        }
      ],
      "source": [
        "jax.device_count('gpu')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LIFmPDWdamno",
        "outputId": "8e21503c-4493-40a1-e68a-b710a4628b19"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[gpu(id=0)]"
            ]
          },
          "metadata": {},
          "execution_count": 18
        }
      ],
      "source": [
        "jax.local_devices()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "sy08TlUVfgY_"
      },
      "outputs": [],
      "source": [
        "arr = jnp.array([1, 42, 31337])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "beWtlcssfl1g",
        "outputId": "c11ffe3e-7ee2-4bc0-f011-8c05ec81a2b7"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "gpu(id=0)"
            ]
          },
          "metadata": {},
          "execution_count": 20
        }
      ],
      "source": [
        "arr.device()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "HQrwov7Zig0t"
      },
      "outputs": [],
      "source": [
        "arr_cpu = jax.device_put(arr, jax.devices('cpu')[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fvhr-gSAivKy",
        "outputId": "82622f22-7f7f-4007-b587-1c258e2964c0"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "CpuDevice(id=0)"
            ]
          },
          "metadata": {},
          "execution_count": 22
        }
      ],
      "source": [
        "arr_cpu.device()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "99hqLCwLin-w",
        "outputId": "d25ae8fe-eeae-4092-ffab-9e7b091b437a"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "gpu(id=0)"
            ]
          },
          "metadata": {},
          "execution_count": 23
        }
      ],
      "source": [
        "arr.device()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "id": "9aQL3DFtjJvm"
      },
      "outputs": [],
      "source": [
        "arr_host = jax.device_get(arr)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "S6ep8EUPjarr",
        "outputId": "94d97bf9-181b-4355-b9a9-7669858a1a28"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "numpy.ndarray"
            ]
          },
          "metadata": {},
          "execution_count": 25
        }
      ],
      "source": [
        "type(arr_host)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7ela_gWXjRde",
        "outputId": "b6f35852-7fc4-4254-9bb1-9e789ead6f89"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([    1,    42, 31337], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 26
        }
      ],
      "source": [
        "arr_host"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8H25pRzgkNnb",
        "outputId": "7e9f7939-2a9e-405f-d89b-16f809e90899"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([    2,    84, 62674], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 27
        }
      ],
      "source": [
        "arr + arr_cpu"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "id": "zMkcgodHknpu"
      },
      "outputs": [],
      "source": [
        "arr_gpu = jax.device_put(arr, jax.devices('gpu')[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gbpuMMGHk1tK",
        "outputId": "afd80fce-0c22-40a8-d7e2-8ee2392bb3c2"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "gpu(id=0)"
            ]
          },
          "metadata": {},
          "execution_count": 29
        }
      ],
      "source": [
        "arr_gpu.device()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xxTzC7Juk3vV",
        "outputId": "4d0636a2-455b-41f3-900d-90749214f601"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Received incompatible devices for jitted computation. Got argument x1 of jax.numpy.add with shape int32[3] and device ids [0] on platform GPU and argument x2 of jax.numpy.add with shape int32[3] and device ids [0] on platform CPU\n"
          ]
        }
      ],
      "source": [
        "try:\n",
        "  arr_gpu + arr_cpu\n",
        "except ValueError as e:\n",
        "  print(e)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MimQQxJVd8OT"
      },
      "source": [
        "## Asyncronous dispatch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "id": "dkMYNMK8d-oZ"
      },
      "outputs": [],
      "source": [
        "import jax"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {
        "id": "8TTzIPOud_c6"
      },
      "outputs": [],
      "source": [
        "a = jnp.array(range(1000000)).reshape((1000,1000))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WL4cK2EgeAOj",
        "outputId": "25e520dd-30ef-41de-9391-a38670ba538a"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(1000, 1000)"
            ]
          },
          "metadata": {},
          "execution_count": 33
        }
      ],
      "source": [
        "a.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {
        "id": "81YXUMwfeUxL",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "6dcb7482-2220-4fde-dd61-03d2661f8898"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "gpu(id=0)"
            ]
          },
          "metadata": {},
          "execution_count": 34
        }
      ],
      "source": [
        "a.device()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%time x = jnp.dot(a,a)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xXmU97i2jB2u",
        "outputId": "93bb84fc-c608-4d64-f6ac-f9e1161ed041"
      },
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 39.9 ms, sys: 1.04 ms, total: 41 ms\n",
            "Wall time: 85.1 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%time x = jnp.dot(a,a).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Q_LOBrs7jc1G",
        "outputId": "269a51f0-1476-4f26-e3a8-a221318c7d19"
      },
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 1.19 ms, sys: 45 µs, total: 1.24 ms\n",
            "Wall time: 6.45 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%time x = np.asarray(jnp.dot(a,a))"
      ],
      "metadata": {
        "id": "9wtnSBNGADQC",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f1aa607e-4a71-4570-d2e4-c330b7119e43"
      },
      "execution_count": 37,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 8.27 ms, sys: 168 µs, total: 8.44 ms\n",
            "Wall time: 7.62 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_cpu = jax.device_put(a, jax.devices('cpu')[0])"
      ],
      "metadata": {
        "id": "vHAlLz9Ck_1O"
      },
      "execution_count": 38,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "a_cpu.device()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "blmetLnalEOq",
        "outputId": "349c8817-96c7-4f09-e8cd-71d7b4d94495"
      },
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "CpuDevice(id=0)"
            ]
          },
          "metadata": {},
          "execution_count": 39
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%time x = jnp.dot(a_cpu,a_cpu).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XSqQAxFRlQ4h",
        "outputId": "9b4d98a6-71f1-4529-9b9f-ef58ecce5686"
      },
      "execution_count": 40,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 333 ms, sys: 527 µs, total: 334 ms\n",
            "Wall time: 192 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mhY8ShWHeBIc"
      },
      "source": [
        "## GPU diagnostics"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 41,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VnA4LuxbTyKb",
        "outputId": "8f2dd07e-ebc6-4dc8-8d18-c68a1d5d7736"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "nvcc: NVIDIA (R) Cuda compiler driver\n",
            "Copyright (c) 2005-2022 NVIDIA Corporation\n",
            "Built on Wed_Sep_21_10:33:58_PDT_2022\n",
            "Cuda compilation tools, release 11.8, V11.8.89\n",
            "Build cuda_11.8.r11.8/compiler.31833905_0\n"
          ]
        }
      ],
      "source": [
        "!nvcc --version"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 42,
      "metadata": {
        "id": "DtPcc6kJUE50"
      },
      "outputs": [],
      "source": [
        "#!pip install --upgrade \"jax[cuda]\" -f https://storage.googleapis.com/jax-releases/jax_releases.html"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 43,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sI5-pu0SUN_n",
        "outputId": "daea3000-4b8f-4c42-9674-e37f77623595"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mon Oct  2 13:02:52 2023       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 525.105.17   Driver Version: 525.105.17   CUDA Version: 12.0     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   52C    P0    28W /  70W |  11743MiB / 15360MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ],
      "source": [
        "gpu_info = !nvidia-smi\n",
        "gpu_info = '\\n'.join(gpu_info)\n",
        "if gpu_info.find('failed') >= 0:\n",
        "  print('Select the Runtime > \"Change runtime type\" menu to enable a GPU accelerator, ')\n",
        "  print('and then re-execute this cell.')\n",
        "else:\n",
        "  print(gpu_info)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Immutability"
      ],
      "metadata": {
        "id": "OzNEJzADTvlC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import jax.numpy as jnp"
      ],
      "metadata": {
        "id": "oKO_DNwDUPSf"
      },
      "execution_count": 44,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp = jnp.array(range(10))\n",
        "a_np  = np.array(range(10))"
      ],
      "metadata": {
        "id": "pHs4TJm1Ty7A"
      },
      "execution_count": 45,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-Eplq9iRKHDM",
        "outputId": "3ac48a01-918f-4467-ddea-bb720cdd7882"
      },
      "execution_count": 46,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 46
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_np[5], a_jnp[5]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GjtKMYxnOqKO",
        "outputId": "4d089b95-00ba-4453-b2ac-e5aea0733c8d"
      },
      "execution_count": 47,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(5, Array(5, dtype=int32))"
            ]
          },
          "metadata": {},
          "execution_count": 47
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_np[5] = 100"
      ],
      "metadata": {
        "id": "9sZC4YCoUUnm"
      },
      "execution_count": 48,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "a_np[5]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HaRvmN0SUWnR",
        "outputId": "561ac73b-40b2-4e64-cb4c-e419d54b2a48"
      },
      "execution_count": 49,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "100"
            ]
          },
          "metadata": {},
          "execution_count": 49
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "try:\n",
        "  a_jnp[5] = 100\n",
        "except TypeError as e:\n",
        "  print(e)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zO8nRaogT3-N",
        "outputId": "423b3fc2-400f-43f5-84d0-5243b60e056e"
      },
      "execution_count": 50,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "'<class 'jaxlib.xla_extension.ArrayImpl'>' object does not support item assignment. JAX arrays are immutable. Instead of ``x[idx] = y``, use ``x = x.at[idx].set(y)`` or another .at[] method: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp = a_jnp.at[5].set(100)"
      ],
      "metadata": {
        "id": "rxAB1h8jT7tb"
      },
      "execution_count": 51,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp[5]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BhSZ0WrfZhsR",
        "outputId": "6fac9383-8996-4098-85cc-8950f52312ab"
      },
      "execution_count": 52,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(100, dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 52
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp[42]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4W91vzNOZmSd",
        "outputId": "e334f000-ceca-4aa0-f2a9-424eb47c2f4d"
      },
      "execution_count": 53,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(9, dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 53
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp.at[42].get()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KmbSWZrNirZd",
        "outputId": "352768a7-77be-4b30-bbfb-7ae97b2ecefc"
      },
      "execution_count": 54,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(9, dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 54
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp.at[42].get(mode='clip')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PFs6_qehi0Hk",
        "outputId": "6d6926e9-aa6c-4854-fe1c-a164486e408f"
      },
      "execution_count": 55,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(9, dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 55
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp.at[42].get(mode='drop')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1F1dy5q7jLW3",
        "outputId": "d7f13888-caaf-4edd-f94d-603cdc181fab"
      },
      "execution_count": 56,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(-2147483648, dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 56
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp.at[42].get(mode='fill', fill_value=-1)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "waFQQytXjcMl",
        "outputId": "064d7a9b-008f-4c92-dab6-a9705900b744"
      },
      "execution_count": 57,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(-1, dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 57
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp = a_jnp.at[42].set(100)\n",
        "a_jnp"
      ],
      "metadata": {
        "id": "6Q-IQhYAjdga",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "3e2ab981-6cf2-4c19-d3ad-4c99d0c607d1"
      },
      "execution_count": 58,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([  0,   1,   2,   3,   4, 100,   6,   7,   8,   9], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 58
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp = a_jnp.at[42].set(100, mode='clip')\n",
        "a_jnp"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6pDDyLIxKkA2",
        "outputId": "0634a817-7c36-4bab-93d9-26475d95b369"
      },
      "execution_count": 59,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([  0,   1,   2,   3,   4, 100,   6,   7,   8, 100], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 59
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Working with float64"
      ],
      "metadata": {
        "id": "ZQoM_CmGOvk5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# this only works on startup!\n",
        "from jax.config import config\n",
        "config.update(\"jax_enable_x64\", True)"
      ],
      "metadata": {
        "id": "Mvz9-cWsOxPC"
      },
      "execution_count": 60,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import jax\n",
        "import jax.numpy as jnp"
      ],
      "metadata": {
        "id": "MtSHfhWyQYFZ"
      },
      "execution_count": 61,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# this will not work on TPU backend. Try using CPU or GPU.\n",
        "x = jnp.array(range(10), dtype=jnp.float64)\n",
        "x.dtype"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BsGsX08bKyFF",
        "outputId": "3660e173-3681-46b0-9a47-0b6bd0e1f7d6"
      },
      "execution_count": 62,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "dtype('float64')"
            ]
          },
          "metadata": {},
          "execution_count": 62
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x.device()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_Jpy24rpO8xb",
        "outputId": "3865f47b-a6a9-467b-f7e0-5fc433e30030"
      },
      "execution_count": 63,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "gpu(id=0)"
            ]
          },
          "metadata": {},
          "execution_count": 63
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "xc = jax.device_put(x, jax.devices('cpu')[0])"
      ],
      "metadata": {
        "id": "5c5u5JTGQ7wP"
      },
      "execution_count": 64,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "xc.device()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3vK1jjnlRCHF",
        "outputId": "14e23c6d-8fea-473f-ada8-5a8040fc4136"
      },
      "execution_count": 65,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "CpuDevice(id=0)"
            ]
          },
          "metadata": {},
          "execution_count": 65
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "xc.dtype"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9bw68etaRHZD",
        "outputId": "f5b8dc1a-dbb8-402a-ae3e-a95a3d2cd038"
      },
      "execution_count": 66,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "dtype('float64')"
            ]
          },
          "metadata": {},
          "execution_count": 66
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "xb16 = jnp.array(range(10), dtype=jnp.bfloat16)\n",
        "xb16.dtype"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DOFosWoRRJY0",
        "outputId": "679ae9e6-b8b3-4977-8719-c778a6bd1ba2"
      },
      "execution_count": 67,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "dtype(bfloat16)"
            ]
          },
          "metadata": {},
          "execution_count": 67
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "xb16.nbytes"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jcY_9IgXSzty",
        "outputId": "45f96941-20c0-4afe-a26b-7086e852f759"
      },
      "execution_count": 68,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "20"
            ]
          },
          "metadata": {},
          "execution_count": 68
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x16 = jnp.array(range(10), dtype=jnp.float16)\n",
        "x16.dtype"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EDxZWsldS4qU",
        "outputId": "d881b0e3-9132-4226-f60a-772a1aec8d70"
      },
      "execution_count": 69,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "dtype('float16')"
            ]
          },
          "metadata": {},
          "execution_count": 69
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x16.nbytes"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_Z4NcluiUouv",
        "outputId": "16348c97-2371-410d-f8d3-ce65fff32dc5"
      },
      "execution_count": 70,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "20"
            ]
          },
          "metadata": {},
          "execution_count": 70
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "xb16+x16"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jC6ynjHBUp4q",
        "outputId": "083b4bd8-c5c8-4ae0-a349-457071922c48"
      },
      "execution_count": 71,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([ 0.,  2.,  4.,  6.,  8., 10., 12., 14., 16., 18.], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 71
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "xb16+xb16"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "d6KaXrhYVget",
        "outputId": "e2317558-83b4-4ff2-f123-19cd67bf1740"
      },
      "execution_count": 72,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([0, 2, 4, 6, 8, 10, 12, 14, 16, 18], dtype=bfloat16)"
            ]
          },
          "metadata": {},
          "execution_count": 72
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## jax.numpy & jax.lax"
      ],
      "metadata": {
        "id": "XzqPR8GS-YzE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "config.update(\"jax_enable_x64\", False)"
      ],
      "metadata": {
        "id": "KdN4dF-mf6bR"
      },
      "execution_count": 84,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import jax.numpy as jnp\n",
        "from jax import lax\n",
        "from jax import random"
      ],
      "metadata": {
        "id": "kNemPL6xVjSK"
      },
      "execution_count": 85,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jnp.add(42, 42.0)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RE4Smt9e-ixs",
        "outputId": "78544f09-fd5c-463e-8263-812ec02ae0dc"
      },
      "execution_count": 86,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(84., dtype=float32, weak_type=True)"
            ]
          },
          "metadata": {},
          "execution_count": 86
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jnp.add(42.0, 42.0)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "z8GOhmXB-mTr",
        "outputId": "7c4cf3f9-614e-41de-e6cf-da1ae9786575"
      },
      "execution_count": 87,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(84., dtype=float32, weak_type=True)"
            ]
          },
          "metadata": {},
          "execution_count": 87
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "try:\n",
        "   lax.add(42, 42.0)\n",
        "except TypeError as e:\n",
        "  print(e)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "G1JzxdU_-rDS",
        "outputId": "7182bc37-7caa-4266-adbc-909ee3224ad2"
      },
      "execution_count": 88,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "lax.add requires arguments to have the same dtypes, got int32, float32. (Tip: jnp.add is a similar function that does automatic type promotion on inputs).\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "lax.add(jnp.float32(42), 42.0)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MTBx3GE6-uNs",
        "outputId": "bd28dd1e-650a-4173-eb06-ac223d9ede82"
      },
      "execution_count": 89,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(84., dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 89
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def random_augmentation(image, augmentations, rng_key):\n",
        "   '''A function that applies a random transformation to an image'''\n",
        "   augmentation_index = random.randint(key=rng_key, minval=0, maxval=len(augmentations), shape=())\n",
        "   augmented_image = lax.switch(augmentation_index, augmentations, image)\n",
        "   return augmented_image"
      ],
      "metadata": {
        "id": "7fyJ9edV-8qO"
      },
      "execution_count": 90,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "add_noise_func = lambda x: x+10\n",
        "horizontal_flip_func = lambda x: x+1\n",
        "rotate_func = lambda x: x+2\n",
        "adjust_colors_func = lambda x: x+3\n",
        "\n",
        "augmentations = [\n",
        "   add_noise_func,\n",
        "   horizontal_flip_func,\n",
        "   rotate_func,\n",
        "   adjust_colors_func\n",
        "]\n"
      ],
      "metadata": {
        "id": "3WJ66AgB_D_o"
      },
      "execution_count": 91,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "image = jnp.array(range(100))"
      ],
      "metadata": {
        "id": "aH-ZJ-DV_Vlu"
      },
      "execution_count": 92,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "random_augmentation(image, augmentations, random.PRNGKey(211))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oRDNt3q1_RJ4",
        "outputId": "0b594b1f-9f4d-4e11-aafb-7ed6a975e527"
      },
      "execution_count": 93,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([  1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,  13,\n",
              "        14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,\n",
              "        27,  28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,  39,\n",
              "        40,  41,  42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,\n",
              "        53,  54,  55,  56,  57,  58,  59,  60,  61,  62,  63,  64,  65,\n",
              "        66,  67,  68,  69,  70,  71,  72,  73,  74,  75,  76,  77,  78,\n",
              "        79,  80,  81,  82,  83,  84,  85,  86,  87,  88,  89,  90,  91,\n",
              "        92,  93,  94,  95,  96,  97,  98,  99, 100], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 93
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "mRahQ4xb_8ma"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}